Spaces:

merasabkuch
/

countThings

Paused

App Files Files Community

dgbkn commited on Jul 27, 2024

Commit

a21ebb4

1 Parent(s): 502613c

added things

Browse files

Files changed (6) hide show

main.py +53 -4
requirements.txt +4 -1
static/index.html +18 -0
upload.html → static/pill_upload.html +15 -1
static/upload.png +0 -0
static/wheat_upload.html +78 -0

main.py CHANGED Viewed

@@ -6,6 +6,12 @@ import cv2
 import numpy as np
 from pillmodel import get_prediction
 import base64
 app = FastAPI()
@@ -17,6 +23,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 @app.post("/predict")
@@ -43,7 +50,49 @@ async def predict(image: UploadFile = File(...)):
     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("upload.html", "r") as file:
-        return HTMLResponse(content=file.read(), status_code=200)

 import numpy as np
 from pillmodel import get_prediction
 import base64
+from fastapi.staticfiles import StaticFiles
+import os
+from inference_sdk import InferenceHTTPClient
 app = FastAPI()
     allow_headers=["*"],
 )
+app.mount("/", StaticFiles(directory="static"), name="static")
 @app.post("/predict")
     return JSONResponse(content={"message": message_to_send, "count": count_dict, "predicted_image": predicted_image_str})
+@app.post("/predict_wheat")
+async def predict_wheat(image: UploadFile = File(...), model_id: str = "grian/1"):
+    contents = await image.read()
+    nparr = np.frombuffer(contents, np.uint8)
+    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+    # delete the image if exists
+    try:
+        os.remove("temp_image.jpg")
+    except:
+        print("temp_image.jpg does not exist")
+    # Save the image to a temporary location
+    temp_image_path = "temp_image.jpg"
+    cv2.imwrite(temp_image_path, img)
+    CLIENT = InferenceHTTPClient(
+    api_url="https://detect.roboflow.com",
+    api_key="PpEebXofNuob5VSx7YP3"
+    )
+    result = CLIENT.infer("temp_image.jpg", model_id=model_id)
+    # Prediction
+    predicted_count = len(result['predictions'])
+    message_to_send = (
+        f"There are {predicted_count} wheat grains."
+    )
+    for prediction in result['predictions']:
+        x = int(prediction['x'])
+        y = int(prediction['y'])
+        width = int(prediction['width'])
+        height = int(prediction['height'])
+        cv2.rectangle(img, (x, y), (x + width, y + height), (0, 255, 0), 2)
+    # Encode predicted image to base64
+    _, buffer = cv2.imencode('.jpg', img)
+    predicted_image_str = base64.b64encode(buffer).decode('utf-8')
+    return JSONResponse(content={"message": message_to_send, "count": predicted_count, "predicted_image": predicted_image_str})

requirements.txt CHANGED Viewed

@@ -9,4 +9,7 @@ shapely
 ultralytics
 fastapi
-uvicorn

 ultralytics
 fastapi
+uvicorn
+inference_sdk
+# roboflow

static/index.html ADDED Viewed

	@@ -0,0 +1,18 @@

+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Detection Models</title>
+    <link href="https://cdn.jsdelivr.net/npm/tailwindcss@2.2.19/dist/tailwind.min.css" rel="stylesheet">
+</head>
+<body class="bg-gray-100 h-screen flex items-center justify-center">
+    <div class="bg-white p-8 rounded shadow-md w-full max-w-sm">
+        <h1 class="text-2xl font-bold mb-4 text-center">Choose a Detection Model</h1>
+        <div class="flex flex-col space-y-4">
+            <a href="/pill_upload.html" class="bg-blue-500 text-white text-center py-2 rounded hover:bg-blue-600">Pill Detection</a>
+            <a href="/wheat_upload.html" class="bg-green-500 text-white text-center py-2 rounded hover:bg-green-600">Wheat Grain Detection</a>
+        </div>
+    </div>
+</body>
+</html>

upload.html → static/pill_upload.html RENAMED Viewed

@@ -9,14 +9,28 @@
         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');

         body {
             font-family: Arial, sans-serif;
         }
+        .clickable-image {
+            cursor: pointer;
+            border: 2px solid #ddd;
+            border-radius: 8px;
+            transition: border-color 0.3s;
+        }
+        .clickable-image:hover {
+            border-color: #007bff;
+        }
     </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/*" capture="environment" class="hidden">
+    <img src="/upload.png" alt="Click to Upload" class="clickable-image mb-4" onclick="triggerFileInput()">
     <button onclick="uploadImage()" class="bg-blue-500 text-white px-4 py-2 rounded">Upload</button>
     <div id="output" class="mt-4"></div>
     <script>
+        function triggerFileInput() {
+            document.getElementById('fileInput').click();
+        }
         async function uploadImage() {
             const fileInput = document.getElementById('fileInput');
             const output = document.getElementById('output');

static/upload.png ADDED Viewed

static/wheat_upload.html ADDED Viewed

	@@ -0,0 +1,78 @@

+<!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;
+        }
+        .clickable-image {
+            cursor: pointer;
+            border: 2px solid #ddd;
+            border-radius: 8px;
+            transition: border-color 0.3s;
+        }
+        .clickable-image:hover {
+            border-color: #007bff;
+        }
+    </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/*" capture="environment" class="hidden">
+    <img src="/upload.png" alt="Click to Upload" class="clickable-image mb-4" onclick="triggerFileInput()">
+    <!-- Dropdown for selecting model_id -->
+    <select id="modelSelect" class="mb-4 px-4 py-2 border rounded">
+        <option value="grian/1">Grian Model</option>
+        <option value="wheat-dataset-new/2">Wheat Dataset New Model</option>
+    </select>
+    <button onclick="uploadImage()" class="bg-blue-500 text-white px-4 py-2 rounded">Upload</button>
+    <div id="output" class="mt-4"></div>
+    <script>
+        function triggerFileInput() {
+            document.getElementById('fileInput').click();
+        }
+        async function uploadImage() {
+            const fileInput = document.getElementById('fileInput');
+            const modelSelect = document.getElementById('modelSelect');
+            const output = document.getElementById('output');
+            const file = fileInput.files[0];
+            const modelId = modelSelect.value;
+            if (!file) {
+                output.innerHTML = 'No image selected.';
+                return;
+            }
+            const formData = new FormData();
+            formData.append('image', file);
+            formData.append('model_id', modelId);  // Append model_id to FormData
+            output.innerHTML = 'Uploading...';
+            try {
+                const response = await fetch('/predict_wheat', {
+                    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>