dnr

best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:504dd2ca20d1e8e7be2d925b840c90c78742e6abe41fb95507dd915d212ca371
+size 92233042

main.py

@@ -0,0 +1,49 @@
+
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse,HTMLResponse
+from fastapi.middleware.cors import CORSMiddleware
+import cv2
+import numpy as np
+from pillmodel import get_prediction
+import base64
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+
+@app.post("/predict")
+async def predict(image: UploadFile = File(...)):
+    contents = await image.read()
+    nparr = np.frombuffer(contents, np.uint8)
+    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+
+    # Save the image to a temporary location
+    # temp_image_path = "temp_image.jpg"
+    # cv2.imwrite(temp_image_path, img)
+
+    # Prediction
+    predicted_image, count_dict = get_prediction(img)
+    # Encode predicted image to base64
+    _, buffer = cv2.imencode('.jpg', predicted_image)
+    predicted_image_str = base64.b64encode(buffer).decode('utf-8')
+
+    # Send a confirmation message
+    message_to_send = (
+        f"There are {count_dict.get('capsules', 0)} capsules and {count_dict.get('tablets', 0)} tablets. "
+        f"A total of {count_dict.get('capsules', 0) + count_dict.get('tablets', 0)} pills."
+    )
+    
+    return JSONResponse(content={"message": message_to_send, "count": count_dict, "predicted_image": predicted_image_str})
+
+@app.get("/", response_class=HTMLResponse)
+async def read_root():
+    with open("index.html", "r") as file:
+        return HTMLResponse(content=file.read(), status_code=200)

pillmodel.py

@@ -0,0 +1,107 @@
+import json
+import numpy as np
+import cv2
+from shapely.geometry import Polygon
+from io import BytesIO
+from ultralytics import YOLO
+import torch
+
+# Define colors
+COLORS = [(98, 231, 4), (228, 161, 0)] # Green and blue
+CLASSES = ['capsules', 'tablets']
+
+
+
+def get_prediction(image):
+    '''
+    Gets image from telebot, make predictions,
+    counts predicted classes, draws dots on image,
+    returns dict with counts and labelled image
+    '''
+    # Load a model
+    model = YOLO('yolov8l.pt')  # load an official model
+    model = YOLO('best.pt')  # load a custom model
+    # image = cv2.imread(image_path)
+
+    # Get prediction
+    prediction = model(image)
+
+    # Get predicted classes
+    predicted_classes = prediction[0].boxes.cls
+    # Get predicted confidence of each class
+    prediction_confidences = prediction[0].boxes.conf
+
+    # Get polygons
+    polygons =  prediction[0].masks.xy
+    # Convert polygons to int32
+    polygons = [polygon.astype(np.int32) for polygon in polygons]
+
+    # Create indices mask that shows what is overlapping polygon has smaller confidence score
+    indices_mask = remove_overlapping_polygons(polygons, prediction_confidences)
+
+    # Create new fixed lists with predicted classes and polygons
+    fixed_predicted_classes = predicted_classes[np.array(indices_mask, dtype=bool)]
+    fixed_polygons = [polygons[i] for i in range(len(indices_mask)) if indices_mask[i] == 1]
+    # fixed_predicted_classes = [predicted_classes[i] for i in range(len(indices_mask)) if indices_mask[i] == 1]
+
+    # Get counts of classes
+    unique, counts = torch.unique(fixed_predicted_classes, return_counts=True)
+    # Get dicts with counts of classes
+    count_dict = {CLASSES[int(key)]: value for key, value in zip(unique.tolist(), counts.tolist())}
+
+    # # Draw polygons
+    # for polygon, predicted_class in zip(fixed_polygons, fixed_predicted_classes):
+    #     cv2.polylines(image, [polygon], True, COLORS[int(predicted_class)])
+
+    # Draw dots
+    for polygon, predicted_class in zip(fixed_polygons, fixed_predicted_classes):
+         # Find center of polygon
+         center_coordinates =  (np.mean(polygon[:, 0], dtype=np.int32), np.mean(polygon[:, 1], dtype=np.int32)) # x and y respectively
+         # Draw a circle
+         cv2.circle(image, center_coordinates, 5, COLORS[int(predicted_class)], 2, cv2.LINE_AA)
+
+    # # Show image with predictions on it
+    # cv2.imshow("Image", image)
+    # cv2.waitKey(0)
+    # cv2.destroyAllWindows()
+    # from google.colab.patches import cv2_imshow
+    # cv2_imshow(image)
+    return image, count_dict
+
+
+def remove_overlapping_polygons(polygons, prediction_confidences):
+    '''
+    Takes polygons, finds overlapping regions,
+    intersection area, overlap percentage,
+    creates indices mask that shows what
+    overlapping polygon has smaller confidence.
+    '''
+
+    # Convert the NumPy arrays to Shapely polygons
+    shapely_polygons = [Polygon(polygon) for polygon in polygons]
+    # Create an empty list with overlapping pairs
+    overlapping_pairs = []
+
+    # Check for overlaps between all pairs of polygons
+    for i in range(len(shapely_polygons)):
+        for j in range(i+1, len(shapely_polygons)):
+            if shapely_polygons[i].intersects(shapely_polygons[j]):
+                # Calculate the percentage of overlap
+                intersection_area = shapely_polygons[i].intersection(shapely_polygons[j]).area
+                overlap_percentage = intersection_area / shapely_polygons[i].area
+                # Add overlapping polygons indexes to list
+                if overlap_percentage > 0.5:
+                    overlapping_pairs.append((i, j))
+
+    # Mask of remains indices
+    indices_mask = [1 for i in range(len(shapely_polygons))]
+
+    # Remove one of the overlapping polygons
+    for first_over_polygon_ind, second_over_polygon_ind in overlapping_pairs:
+        # Find index that has the smallest prediction confidence
+        first_has_bigger_conf = prediction_confidences[first_over_polygon_ind] >= prediction_confidences[second_over_polygon_ind]
+        index_small_conf = [first_over_polygon_ind, second_over_polygon_ind][first_has_bigger_conf]
+        # Set value with smaller confidence to 0 in indices_mask
+        indices_mask[index_small_conf] = 0
+
+    return indices_mask

requirements.txt

@@ -0,0 +1,9 @@
+cvzone
+ultralytics
+numpy
+pandas
+torch
+
+opencv-python
+shapely
+ultralytics

upload.html

@@ -0,0 +1,52 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Image Upload</title>
+    <link href="https://cdn.jsdelivr.net/npm/tailwindcss@2.2.19/dist/tailwind.min.css" rel="stylesheet">
+    <style>
+        body {
+            font-family: Arial, sans-serif;
+        }
+    </style>
+</head>
+<body class="flex flex-col items-center justify-center h-screen bg-gray-100">
+    <h1 class="text-2xl font-bold mb-4">Upload an Image</h1>
+    <input type="file" id="fileInput" accept="image/*" class="mb-4">
+    <button onclick="uploadImage()" class="bg-blue-500 text-white px-4 py-2 rounded">Upload</button>
+    <div id="output" class="mt-4"></div>
+    <script>
+        async function uploadImage() {
+            const fileInput = document.getElementById('fileInput');
+            const output = document.getElementById('output');
+            const file = fileInput.files[0];
+            if (!file) {
+                output.innerHTML = 'No image selected.';
+                return;
+            }
+
+            const formData = new FormData();
+            formData.append('image', file);
+
+            output.innerHTML = 'Uploading...';
+
+            try {
+                const response = await fetch('/predict', {
+                    method: 'POST',
+                    body: formData
+                });
+                const result = await response.json();
+                const predictedImageSrc = `data:image/jpeg;base64,${result.predicted_image}`;
+                output.innerHTML = `
+                    <p>${result.message}</p>
+                    <img src="${predictedImageSrc}" alt="Predicted Image" class="mt-4">
+                `;
+            } catch (error) {
+                output.innerHTML = 'Failed to get prediction';
+                console.error(error);
+            }
+        }
+    </script>
+</body>
+</html>