feat: add object detection

main.py

@@ -3,11 +3,17 @@ from fastapi import FastAPI, File, UploadFile
 import tensorflow as tf
 from PIL import Image
 from io import BytesIO
+from ultralytics import YOLO
+import cv2
+from datetime import datetime
+from fastapi.responses import FileResponse
 
 app = FastAPI()
 
 labels = []
-model = tf.keras.models.load_model('./models.h5')
+classification_model = tf.keras.models.load_model('./models.h5')
+detection_model = YOLO('./best.pt')
+
 with open("labels.txt") as f:
     for line in f:
         labels.append(line.replace('\n', ''))
@@ -17,13 +23,116 @@ def classify_image(img):
     img_array = np.asarray(img.resize((224, 224)))[..., :3]
     img_array = img_array.reshape((1, 224, 224, 3))  # Add batch dimension
     img_array = tf.keras.applications.efficientnet.preprocess_input(img_array)
-    prediction = model.predict(img_array).flatten()
+    prediction = classification_model.predict(img_array).flatten()
     confidences = {labels[i]: float(prediction[i]) for i in range(90)}
     # Sort the confidences dictionary by value and get the top 3 items
     # top_3_confidences = dict(sorted(confidences.items(), key=lambda item: item[1], reverse=True)[:3])
     
     return confidences
 
+def animal_detect_and_classify(img_path):
+    # Read the image
+    img = cv2.imread(img_path)
+
+    # Pass the image through the detection model and get the result
+    detect_results = detection_model(img)
+
+    combined_results = []
+    # print("dss", detect_results[0])
+    # Iterate over the detected objects
+    # Iterate over detections
+    for result in detect_results:
+        for box in result.boxes:
+            # print(box)
+            # Crop the RoI
+            x1, y1, x2, y2 = map(int, box.xyxy[0])
+            detect_img = img[y1:y2, x1:x2]
+            # Convert the image to RGB format
+            detect_img = cv2.cvtColor(detect_img, cv2.COLOR_BGR2RGB)
+
+            # Resize the input image to the expected shape (224, 224)
+            detect_img = cv2.resize(detect_img, (224, 224))
+
+            # Convert the image to a numpy array
+            inp_array = np.array(detect_img)
+
+            # Reshape the array to match the expected input shape
+            inp_array = inp_array.reshape((-1, 224, 224, 3))
+
+            # Preprocess the input array
+            inp_array = tf.keras.applications.efficientnet.preprocess_input(inp_array)
+
+            # Make predictions using the classification model
+            prediction = classification_model.predict(inp_array)
+            # Map predictions to labels
+            threshold = 0.75
+            predicted_labels = [labels[np.argmax(pred)] if np.max(pred) >= threshold else "animal" for pred in prediction]
+            print(predicted_labels)
+            combined_results.append(((x1, y1, x2, y2), predicted_labels))
+
+    return combined_results
+
+def generate_color(class_name):
+    # Generate a hash from the class name
+    color_hash = hash(class_name)
+    print(color_hash)
+    # Normalize the hash value to fit within the range of valid color values (0-255)
+    color_hash = abs(color_hash) % 16777216
+    R = color_hash//(256*256)
+    G = (color_hash//256) % 256
+    B = color_hash % 256
+    # Convert the hash value to RGB color format
+    color = (R, G, B)
+
+    return color
+
+def plot_detected_rectangles(image, detections, output_path):
+    # Create a copy of the image to draw on
+    img_with_rectangles = image.copy()
+
+    # Iterate over each detected rectangle and its corresponding class name
+    for rectangle, class_names in detections:
+        # Extract the coordinates of the rectangle
+        x1, y1, x2, y2 = rectangle
+
+        # Generate a random color
+        color = generate_color(class_names[0])
+
+        # Draw the rectangle on the image
+        cv2.rectangle(img_with_rectangles, (x1, y1), (x2, y2), color, 2)
+
+        # Put the class names above the rectangle
+        for i, class_name in enumerate(class_names):
+            cv2.putText(img_with_rectangles, class_name, (x1, y1 - 10 - i*20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+    # Show the image with rectangles and class names
+    cv2.imwrite(output_path, img_with_rectangles)
+
+
+# Call the animal_detect_and_classify function to get detections
+detections = animal_detect_and_classify('/content/cat_tiger.jpg')
+
+# Plot the detected rectangles with their corresponding class names
+plot_detected_rectangles(cv2.imread('/content/cat_tiger.jpg'), detections)
+
+
+@app.post("/predict/v2")
+async def predict_v2(file: UploadFile = File(...)):
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = timestamp + file.filename
+    contents = await file.read()
+    image = Image.open(BytesIO(contents))
+    image.save("input/" + filename)
+    detections = animal_detect_and_classify("input/" + filename)
+    plot_detected_rectangles(cv2.imread("input/" + filename), detections, "output/" + filename)
+    return {"message": "Detection and classification completed successfully"}
+
+@app.get("/image/")
+async def get_image(image_name: str):
+    # Assume the images are stored in a directory named "images"
+    image_path = f"images/{image_name}"
+    return FileResponse(image_path)
+
 @app.post("/predict")
 async def predict(file: bytes = File(...)):
     img = Image.open(BytesIO(file))

requirements.txt

@@ -10,4 +10,7 @@ uvicorn
 python-multipart
 numpy==1.25.2
 Pillow==9.4.0
-keras==2.15.0
+keras==2.15.0
+ultralytics
+squarify
+cv2