Update app.py

app.py

@@ -1,5 +1,5 @@
 from fastapi import FastAPI, File, Form, UploadFile
-from fastapi.responses import StreamingResponse
+from fastapi.responses import JSONResponse
 import torch
 from transformers import GroundingDinoForObjectDetection, AutoProcessor
 from PIL import Image, ImageDraw
@@ -10,13 +10,13 @@ import threading
 
 app = FastAPI()
 
-# MySQL database connection setting
+# MySQL database connection settings
 username = 'ukrqsqxg_kacafix'
 password = 'kdm#{k&4@&Y+'
 host = 'www.kacafix.com'
 database = 'ukrqsqxg_millbox_storage'
 
-# Create a connection to the MySQL databased
+# Create a connection to the MySQL database
 cnx = mysql.connector.connect(
     user=username,
     password=password,
@@ -32,10 +32,6 @@ model = GroundingDinoForObjectDetection.from_pretrained('IDEA-Research/grounding
 processor = AutoProcessor.from_pretrained("IDEA-Research/grounding-dino-tiny")
 model.to(device)
 
-# Cache the pre-trained model and processor
-model_cache = model
-processor_cache = processor
-
 lock = threading.Lock()
 
 @app.post("/predict")
@@ -52,13 +48,12 @@ async def predict(
     labels = [label if label.endswith(".") else label + "." for label in labels]
     labels = " ".join(labels)
 
-    # Use cached model and processor
     with lock:
-        inputs = processor_cache(images=image, text=labels, return_tensors="pt").to(device)
+        inputs = processor(images=image, text=labels, return_tensors="pt").to(device)
         with torch.no_grad():
-            outputs = model_cache(**inputs)
+            outputs = model(**inputs)
 
-    result = processor_cache.post_process_grounded_object_detection(
+    result = processor.post_process_grounded_object_detection(
         outputs,
         inputs.input_ids,
         box_threshold=box_threshold,
@@ -81,7 +76,8 @@ async def predict(
     # Save the data in the MySQL database asynchronously
     asyncio.create_task(save_to_database(labels, box_threshold, text_threshold, output_image_io.getvalue()))
 
-    return JSONResponse(content={"message": "complete"})
+    # Return a success message instead of the image
+    return JSONResponse(content={"message": "Scan Completed Successfully"})
 
 async def save_to_database(labels, box_threshold, text_threshold, output_image):
     query = ("INSERT INTO model_requests (labels, box_threshold, text_threshold, output_image) "