Update app.py

app.py

@@ -13,7 +13,7 @@ app = FastAPI()
 
 # Load SmolVLM-Instruct model and processor
 model_id = "HuggingFaceTB/SmolVLM-Instruct"
-processor = AutoProcessor.from_pretrained(model_id, token=os.environ.get("HF_TOKEN"))
+processor = AutoProcessor.from_pretrained(model_id, token=os.environ.get("HF_TOKEN"), padding=True)
 model = AutoModelForVision2Seq.from_pretrained(model_id, token=os.environ.get("HF_TOKEN"))
 
 # Harmful objects list for detection
@@ -31,23 +31,26 @@ async def predict(files: List[UploadFile] = File(...)):
             image = Image.open(io.BytesIO(image_data)).convert("RGB")
             
             # Generate description
-            inputs = processor(text=["<image> Describe the image in detail."], images=[image], return_tensors="pt")
-            outputs = model.generate(**inputs, max_length=100)
-            description = processor.decode(outputs[0], skip_special_tokens=True).replace("Describe the image in detail.", "").strip()
+            prompt_desc = "<image> Provide a detailed description of the image, including objects, colors, and context."
+            inputs = processor(text=[prompt_desc], images=[image], return_tensors="pt", padding=True)
+            outputs = model.generate(**inputs, max_length=512)
+            description = processor.decode(outputs[0], skip_special_tokens=True).replace(prompt_desc, "").strip()
             
             # Extract signs/number plates (OCR)
-            inputs_ocr = processor(text=["<image> Extract all text visible in the image."], images=[image], return_tensors="pt")
-            ocr_outputs = model.generate(**inputs_ocr, max_length=100)
-            signs_text = processor.decode(ocr_outputs[0], skip_special_tokens=True).replace("Extract all text visible in the image.", "").strip()
+            prompt_ocr = "<image> Extract all visible text in the image, such as signs or license plates."
+            inputs_ocr = processor(text=[prompt_ocr], images=[image], return_tensors="pt", padding=True)
+            ocr_outputs = model.generate(**inputs_ocr, max_length=512)
+            signs_text = processor.decode(ocr_outputs[0], skip_special_tokens=True).replace(prompt_ocr, "").strip()
             
             # Detect harmful objects/blood
-            inputs_detect = processor(text=["<image> List all objects in the image."], images=[image], return_tensors="pt")
-            detect_outputs = model.generate(**inputs_detect, max_length=100)
+            prompt_detect = "<image> Identify any harmful objects (e.g., knife, gun, blood, syringe, bomb, blade) in the image. List them explicitly."
+            inputs_detect = processor(text=[prompt_detect], images=[image], return_tensors="pt", padding=True)
+            detect_outputs = model.generate(**inputs_detect, max_length=512)
             detected_objects = processor.decode(detect_outputs[0], skip_special_tokens=True).lower()
             harmful_detected = any(obj in detected_objects for obj in harmful_objects) and "Detected: " + ", ".join([obj for obj in harmful_objects if obj in detected_objects]) or "None detected"
             
             # Get image embedding for similarity
-            inputs_emb = processor(images=[image], return_tensors="pt")
+            inputs_emb = processor(images=[image], return_tensors="pt", padding=True)
             with torch.no_grad():
                 emb = model.vision_model(inputs_emb["pixel_values"]).last_hidden_state.mean(dim=1).cpu().numpy()
             image_embeddings.append(emb)
@@ -89,23 +92,26 @@ def gradio_predict(*images):
             image = Image.fromarray(image).convert("RGB")
             
             # Generate description
-            inputs = processor(text=["<image> Describe the image in detail."], images=[image], return_tensors="pt")
-            outputs = model.generate(**inputs, max_length=100)
-            description = processor.decode(outputs[0], skip_special_tokens=True).replace("Describe the image in detail.", "").strip()
+            prompt_desc = "<image> Provide a detailed description of the image, including objects, colors, and context."
+            inputs = processor(text=[prompt_desc], images=[image], return_tensors="pt", padding=True)
+            outputs = model.generate(**inputs, max_length=512)
+            description = processor.decode(outputs[0], skip_special_tokens=True).replace(prompt_desc, "").strip()
             
             # Extract signs/number plates (OCR)
-            inputs_ocr = processor(text=["<image> Extract all text visible in the image."], images=[image], return_tensors="pt")
-            ocr_outputs = model.generate(**inputs_ocr, max_length=100)
-            signs_text = processor.decode(ocr_outputs[0], skip_special_tokens=True).replace("Extract all text visible in the image.", "").strip()
+            prompt_ocr = "<image> Extract all visible text in the image, such as signs or license plates."
+            inputs_ocr = processor(text=[prompt_ocr], images=[image], return_tensors="pt", padding=True)
+            ocr_outputs = model.generate(**inputs_ocr, max_length=512)
+            signs_text = processor.decode(ocr_outputs[0], skip_special_tokens=True).replace(prompt_ocr, "").strip()
             
             # Detect harmful objects/blood
-            inputs_detect = processor(text=["<image> List all objects in the image."], images=[image], return_tensors="pt")
-            detect_outputs = model.generate(**inputs_detect, max_length=100)
+            prompt_detect = "<image> Identify any harmful objects (e.g., knife, gun, blood, syringe, bomb, blade) in the image. List them explicitly."
+            inputs_detect = processor(text=[prompt_detect], images=[image], return_tensors="pt", padding=True)
+            detect_outputs = model.generate(**inputs_detect, max_length=512)
             detected_objects = processor.decode(detect_outputs[0], skip_special_tokens=True).lower()
             harmful_detected = any(obj in detected_objects for obj in harmful_objects) and "Detected: " + ", ".join([obj for obj in harmful_objects if obj in detected_objects]) or "None detected"
             
             # Get image embedding for similarity
-            inputs_emb = processor(images=[image], return_tensors="pt")
+            inputs_emb = processor(images=[image], return_tensors="pt", padding=True)
             with torch.no_grad():
                 emb = model.vision_model(inputs_emb["pixel_values"]).last_hidden_state.mean(dim=1).cpu().numpy()
             image_embeddings.append(emb)
@@ -153,7 +159,7 @@ iface = gr.Interface(
     inputs=[gr.Image(label=f"Upload Image {i+1}") for i in range(3)],  # Allow up to 3 images
     outputs=gr.Textbox(label="Results"),
     title="VisionSage: Image Analysis with SmolVLM",
-    description="Upload up to 3 images to get descriptions, extract signs/number plates, detect harmful objects/blood, and compute similarity to the first image."
+    description="Upload up to 3 images to get detailed descriptions, extract signs/number plates, detect harmful objects/blood, and compute similarity to the first image."
 )
 
 if __name__ == "__main__":