Update app.py

app.py

@@ -78,40 +78,14 @@ def generate_gemini_response(input_prompt, image):
 
 # Object detection part
 def detect_objects(image):
-    # Ensure the image is in RGB mode
-    if image.mode != "RGB":
-        image = image.convert("RGB")
-    
-    # Convert the image to a numpy array
-    image_np = np.array(image)
-    
-    # Load the processor and model
     processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
     model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
-    
-    # Process the image for detection
-    inputs = processor(images=image_np, return_tensors="pt")
+
+    inputs = processor(images=image, return_tensors="pt")
+
     outputs = model(**inputs)
 
     target_sizes = torch.tensor([image.size[::-1]])
-    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
-
-    detected_cars = []
-    draw = ImageDraw.Draw(image)
-
-    # Loop through detections and filter only "car" class (ID 3 for COCO dataset)
-    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-        if model.config.id2label[label.item()] == 'car' and score.item() > 0.9:
-            box = [round(i, 2) for i in box.tolist()]
-            # Crop the detected car
-            cropped_car = image.crop(box)
-            # Convert the cropped image to bytes
-            cropped_car_bytes = image_to_bytes(cropped_car)
-            detected_cars.append((cropped_car_bytes, box))
-
-            # Draw bounding box around the car
-            draw.rectangle(box, outline="red", width=3)
-            draw.text((box[0], box[1]), f"Car: {round(score.item(), 2)}", fill="red")
     
     return image, detected_cars