Update app.py

app.py

@@ -17,6 +17,10 @@ app = FastAPI()
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+
+class_mapping = {'tb': 0, 'healthy': 1, 'sick_but_no_tb': 2}
+reverse_mapping = {v: k for k, v in class_mapping.items()}
+
 def load_model():
     # config = read_params(config_path)
     model = timm.create_model('convnext_base.clip_laiona', pretrained=True, num_classes=3)
@@ -52,27 +56,56 @@ def allowed_file(filename):
     return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
 
 
-@app.get("/predict")
-async def predict(file: UploadFile = File(...)):
-    """
-    Perform prediction on the uploaded image
-    """
+# @app.get("/predict")
+# async def predict(file: UploadFile = File(...)):
+#     """
+#     Perform prediction on the uploaded image
+#     """
 
-    logger.info('API predict called')
+#     logger.info('API predict called')
 
-    if not allowed_file(file.filename):
-        raise HTTPException(status_code=400, detail="Format not supported")
+#     if not allowed_file(file.filename):
+#         raise HTTPException(status_code=400, detail="Format not supported")
     
-    try:
-        img_bytes = await file.read()
-        class_name = get_prediction(img_bytes)
-        logger.info(f'Prediction: {class_name}')
-        return JSONResponse(content={"class_name": class_name})
-    except Exception as e:
-        logger.error(f'Error: {str(e)}')
-        return JSONResponse(content={"error": str(e), "trace": traceback.format_exc()}, status_code=500)
+#     try:
+#         img_bytes = await file.read()
+#         class_name = get_prediction(img_bytes)
+#         logger.info(f'Prediction: {class_name}')
+#         return JSONResponse(content={"class_name": class_name})
+#     except Exception as e:
+#         logger.error(f'Error: {str(e)}')
+#         return JSONResponse(content={"error": str(e), "trace": traceback.format_exc()}, status_code=500)
      
 
-# @app.get("/")
-# def greet_json():
-#     return {"Hello": "World!"}
+# # @app.get("/")
+# # def greet_json():
+# #     return {"Hello": "World!"}
+
+
+import torch
+import requests
+from PIL import Image
+from torchvision import transforms
+
+# model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True).eval()
+
+
+# Download human-readable labels for ImageNet.
+# response = requests.get("https://git.io/JJkYN")
+# labels = response.text.split("\n")
+
+
+def predict(inp):
+  inp = transforms.ToTensor()(inp).unsqueeze(0)
+  with torch.no_grad():
+    prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
+    # confidences = {labels[i]: float(prediction[i]) for i in range(1000)}
+    prediction = reverse_mapping[prediction]
+  return prediction 
+
+
+import gradio as gr
+
+gr.Interface(fn=predict,
+             inputs=gr.Image(type="pil"),
+             outputs=gr.Label(num_top_classes=3)).launch()