Update app.py

app.py

@@ -1,47 +1,78 @@
+import os
 import torch
+from transformers import AutoModelForCausalLM, AutoProcessor
 from PIL import Image
-from transformers import AutoProcessor, AutoModelForVision2Seq
-import gradio as gr
-import os
-
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-token = os.environ.get("HF_TOKEN")
-
-processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM-Instruct", token=token)
-model = AutoModelForVision2Seq.from_pretrained(
-    "HuggingFaceTB/SmolVLM-Instruct",
-    torch_dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
-    _attn_implementation="flash_attention_2" if DEVICE == "cuda" else "eager",
-    token=token
-).to(DEVICE)
-
-def describe_image(image):
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image"},
-                {"type": "text", "text": "Describe this image in detail."}
-            ]
-        },
-    ]
+import numpy as np
+from fastapi import FastAPI, UploadFile, File
+from typing import List
+import io
+import uvicorn
+from pyngrok import ngrok
 
-    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
-    inputs = processor(text=prompt, images=[image], return_tensors="pt").to(DEVICE)
+# Initialize FastAPI app
+app = FastAPI()
 
-    with torch.no_grad():
-        generated_ids = model.generate(**inputs, max_new_tokens=500)
+# Load SmolVLM-Instruct model and processor
+model_id = "HuggingFaceTB/SmolVLM-Instruct"
+processor = AutoProcessor.from_pretrained(model_id, token=os.environ.get("HF_TOKEN"))
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16).to("cuda")
 
-    result = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    return result
+# Harmful objects list for detection
+harmful_objects = ["knife", "gun", "weapon", "blood", "syringe"]
 
-demo = gr.Interface(
-    fn=describe_image,
-    inputs=gr.Image(type="pil", label="Upload an Image"),
-    outputs="text",
-    title="VISIONSAGE",
-    description="Upload an image and get a detailed description."
-)
+@app.post("/predict")
+async def predict(files: List[UploadFile] = File(...)):
+    results = []
+    image_embeddings = []
+    
+    for file in files:
+        # Read image
+        image_data = await file.read()
+        image = Image.open(io.BytesIO(image_data)).convert("RGB")
+        
+        # Generate description
+        inputs = processor(text="Describe the image in detail.", images=image, return_tensors="pt").to("cuda")
+        outputs = model.generate(**inputs, max_length=100)
+        description = processor.decode(outputs[0], skip_special_tokens=True).replace("Describe the image in detail.", "").strip()
+        
+        # Extract signs/number plates (OCR)
+        inputs_ocr = processor(text="Extract all text visible in the image.", images=image, return_tensors="pt").to("cuda")
+        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()
+        
+        # Detect harmful objects/blood
+        inputs_detect = processor(text="List all objects in the image.", images=image, return_tensors="pt").to("cuda")
+        detect_outputs = model.generate(**inputs_detect, max_length=100)
+        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" or "None detected"
+        
+        # Get image embedding for similarity
+        inputs_emb = processor(images=image, return_tensors="pt").to("cuda")
+        with torch.no_grad():
+            emb = model.vision_tower(inputs_emb["pixel_values"]).last_hidden_state.mean(dim=1).cpu().numpy()
+        image_embeddings.append(emb)
+        
+        results.append({
+            "description": description if description else "No description generated.",
+            "signs": signs_text if signs_text else "None detected",
+            "harmful": harmful_detected
+        })
+    
+    # Compute similarity scores (cosine similarity to first image)
+    if len(image_embeddings) > 1:
+        base_embedding = image_embeddings[0]
+        for i in range(1, len(image_embeddings)):
+            sim = np.dot(base_embedding, image_embeddings[i].T) / (
+                np.linalg.norm(base_embedding) * np.linalg.norm(image_embeddings[i])
+            )
+            results[i]["similarity"] = float(sim[0][0])
+    
+    return results
 
-demo.launch()
+if __name__ == "__main__":
+    # Start ngrok tunnel
+    public_url = ngrok.connect(8000).public_url
+    print(f"Public URL: {public_url}")
+    
+    # Run FastAPI server
+    uvicorn.run(app, host="0.0.0.0", port=8000)