Test deployment

.gitignore

@@ -0,0 +1,4 @@
+.venv
+__pycache__
+*.pyc
+*.gradio

classification_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:74a6325f0e12a18554e204106551536a00fafe003294deb6836fc09082d2a8ee
+size 32390942

classification_model.py

@@ -0,0 +1,70 @@
+import io
+import os
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+from PIL import Image
+import timm
+
+# Define your pill classes (ensure this matches your training setup)
+PILL_CLASSES = {
+    'acc': 0, 'advil': 1, 'akineton': 2, 'algoflex': 3, 'algopyrin': 4, 'ambroxol': 5, 
+    'apranax': 6, 'aspirin': 7, 'atoris': 8, 'atorvastatin': 9, 'betaloc': 10, 
+    'bila': 11, 'c': 12, 'calci': 13, 'cataflam': 14, 'cetirizin': 15, 'co': 16, 
+    'cold': 17, 'coldrex': 18, 'concor': 19, 'condrosulf': 20, 'controloc': 21, 
+    'covercard': 22, 'coverex': 23, 'diclopram': 24, 'donalgin': 25, 'dorithricin': 26, 
+    'doxazosin': 27, 'dulodet': 28, 'dulsevia': 29, 'enterol': 30, 'escitil': 31, 
+    'favipiravir': 32, 'frontin': 33, 'furon': 34, 'ibumax': 35, 'indastad': 36, 
+    'jutavit': 37, 'kalcium': 38, 'kalium': 39, 'ketodex': 40, 'koleszterin': 41, 
+    'l': 42, 'lactamed': 43, 'lactiv': 44, 'laresin': 45, 'letrox': 46, 'lordestin': 47, 
+    'magne': 48, 'mebucain': 49, 'merckformin': 50, 'meridian': 51, 'metothyrin': 52, 
+    'mezym': 53, 'milgamma': 54, 'milurit': 55, 'naprosyn': 56, 'narva': 57, 
+    'naturland': 58, 'nebivolol': 59, 'neo': 60, 'no': 61, 'noclaud': 62, 
+    'nolpaza': 63, 'nootropil': 64, 'normodipine': 65, 'novo': 66, 'nurofen': 67, 
+    'ocutein': 68, 'olicard': 69, 'panangin': 70, 'pantoprazol': 71, 'provera': 72, 
+    'quamatel': 73, 'reasec': 74, 'revicet': 75, 'rhinathiol': 76, 'rubophen': 77, 
+    'salazopyrin': 78, 'sedatif': 79, 'semicillin': 80, 'sicor': 81, 'sinupret': 82, 
+    'sirdalud': 83, 'strepfen': 84, 'strepsils': 85, 'syncumar': 86, 'teva': 87, 
+    'theospirex': 88, 'tricovel': 89, 'tritace': 90, 'urotrin': 91, 'urzinol': 92, 
+    'valeriana': 93, 'verospiron': 94, 'vita': 95, 'vitamin': 96, 'voltaren': 97, 
+    'xeter': 98, 'zadex': 99
+}
+
+# Set device to CPU
+device = torch.device("cpu")
+
+# Instantiate the model architecture (same as training)
+model = timm.create_model("rexnet_150", pretrained=True, num_classes=len(PILL_CLASSES))
+model.to(device)
+
+# Load the trained state dict
+model_path = os.path.join("classification_model.pth")
+model.load_state_dict(torch.load(model_path, map_location=device))
+model.eval()
+
+# Define image transformations
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+def classify_medicine(image_bytes):
+    """Convert image bytes to prediction using the loaded model."""
+    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    input_tensor = transform(image).unsqueeze(0).to(device)
+    
+    with torch.no_grad():
+        output = model(input_tensor)
+    
+    probabilities = F.softmax(output[0], dim=0)
+    class_index = torch.argmax(probabilities).item()
+    confidence = probabilities[class_index].item()
+    
+    # Invert the PILL_CLASSES dictionary for easy lookup
+    PILL_CLASSES_INVERTED = {v: k for k, v in PILL_CLASSES.items()}
+    pill_class = PILL_CLASSES_INVERTED.get(class_index, "Unknown")
+    
+    return {"class_index": class_index, "pill_class": pill_class, "confidence": confidence}
+
+export = classify_medicine

main.py

@@ -0,0 +1,31 @@
+import gradio as gr
+from classification_model import classify_medicine
+from ocr_model import perform_ocr
+
+def classify_image(image):
+    with open(image, "rb") as f:
+        image_bytes = f.read()
+    result = classify_medicine(image_bytes)
+    return f"Class: {result['pill_class']}, Confidence: {result['confidence']:.2f}"
+
+def ocr_image(image):
+    with open(image, "rb") as f:
+        image_bytes = f.read()
+    return perform_ocr(image_bytes)
+
+with gr.Blocks() as app:
+    gr.Markdown("## Medicine Classification and OCR App")
+    
+    with gr.Tab("Classify Medicine"):
+        image_input = gr.Image(type="filepath")
+        classify_button = gr.Button("Classify")
+        output_text = gr.Textbox()
+        classify_button.click(classify_image, inputs=image_input, outputs=output_text)
+    
+    with gr.Tab("OCR Extraction"):
+        image_input_ocr = gr.Image(type="filepath")
+        ocr_button = gr.Button("Extract Text")
+        ocr_output = gr.Textbox()
+        ocr_button.click(ocr_image, inputs=image_input_ocr, outputs=ocr_output)
+
+app.launch()

ocr_model.py

@@ -0,0 +1,36 @@
+import io
+import requests
+from PIL import Image
+
+def perform_ocr(image_bytes):
+    if not image_bytes:
+        raise ValueError("Empty image bytes provided")
+    # Validate image bytes
+    try:
+        Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    except Exception as e:
+        raise ValueError(f"Invalid image bytes provided: {e}")
+
+    # OCR.space API endpoint and payload (using the free 'helloworld' key)
+    api_url = "https://api.ocr.space/parse/image"
+    payload = {
+        'apikey': 'helloworld',  # Free API key with usage limits
+        'language': 'eng'
+    }
+    files = {
+        'file': ('image.jpg', image_bytes)
+    }
+
+    response = requests.post(api_url, data=payload, files=files)
+    result = response.json()
+
+    if result.get("IsErroredOnProcessing"):
+        error = result.get("ErrorMessage") or "Unknown error"
+        raise ValueError(f"OCR processing error: {error}")
+
+    parsed_text = result.get("ParsedResults")[0].get("ParsedText", "")
+    paragraphs = parsed_text.split('\n')
+    formatted_text = "\n\n".join(p.strip() for p in paragraphs if p.strip())
+    return formatted_text
+
+export = perform_ocr

requirements.txt

@@ -0,0 +1,7 @@
+gradio
+transformers  # For Hugging Face models
+torch
+torchvision
+pillow        # For image processing
+python-dotenv # For loading environment variables
+timm