Team Ai
Apppublic

Pranavpai0309/Image_Processing

sourceHugging Faceupdated 1y agoView on Hugging Face
0likes
ModelCode.py277 linesDownload Raw Back to root
1import torch2import clip3from PIL import Image4import gc5from transformers import BlipProcessor, BlipForConditionalGeneration6import pytesseract7 8device = "cpu"9 10clip_model = None11clip_preprocess = None12blip_processor = None13blip_model = None14 15blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")16blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)17 18MAX_SIZE = (1024, 1024)19 20def load_and_resize_image(image_path):21    img = Image.open(image_path)22    if img.size[0] * img.size[1] > 89478485:23        print("Image too large, resizing...")24        img.thumbnail(MAX_SIZE)25    return img26 27 28ayurvedic_prompts = [29    "Ayurvedic medicine",30    "Ayurvedic treatment center",31    "Panchakarma therapy",32    "Natural herbal remedies",33    "Traditional Indian healing",34    "Herbal powder and oils",35    "Ayurvedic massage",36    "Siddha or Ayurveda practice",37    "Holy basil and turmeric",38    "Bowl of dried medicinal herbs",39    "Ayurvedic detox therapy",40    "Spiritual yoga and healing",41    "Ashwagandha root and powder",42    "Boiling herbal concoction",43    "Ayurvedic spa setting",44    "Handmade herbal soap",45    "Mortar and pestle with herbs",46    "Ayurvedic doctor in traditional dress",47    "Ayurvedic clinic interior",48    "Holy plants in a healing ritual",49    "Neem and tulsi leaves",50    "Ayurveda-inspired diet plan",51    "Copper water vessel with lemon",52    "Meditation with Ayurvedic setup",53    "Natural treatment using roots and leaves",54    "Herbal tea with cumin and ginger",55    "Ayurvedic facial therapy",56    "Healing stones with herbs",57    "Triphala powder and fruit",58    "Organic herbal tinctures",59    "Traditional ayurvedic scriptures",60    "Dry herbs in glass jars",61    "Oil pulling with sesame oil",62    "Chyawanprash herbal jam",63    "Ayurvedic nutrition consultation",64    "Cumin water detox",65    "Shatavari plant remedy",66    "Ginger and turmeric decoction",67    "Ayurveda in daily routine",68    "Herbal smoke cleansing",69    "Medicinal herb garden",70    "Ayurveda lifestyle book",71    "Ayurvedic pulse diagnosis",72    "Bottle of Ayurvedic syrup",73    "Marma therapy session",74    "Amla powder for hair and skin",75    "Natural immunity boosters",76    "Ayurvedic golden milk",77    "Cleansing herbal juices",78    "Ayurvedic retreat",79    "Ayurvedic foot bath",80    "Mulethi sticks for cough",81    "Herbal incense in meditation",82    "Traditional Kerala Ayurveda",83    "Kapha balancing diet",84    "Herbal body wrap",85    "Sandalwood paste application",86    "Turmeric face mask",87    "Ayurvedic diet for digestion",88    "Herbal ghee preparation",89    "Ayurvedic healing altar",90    "Handmade Ayurvedic candles",91    "Herbal steam therapy",92    "Ayurvedic tonics in bottles",93    "Ayurveda consultation with practitioner",94    "Dhoop sticks and herbs",95    "Ayurveda dosha chart",96    "Coriander and fennel tea",97    "Ayurvedic herb grinder",98    "Herbal supplement capsules",99    "Boiling neem leaves",100    "Yoga and Ayurveda combo session",101    "Sun-dried medicinal roots",102    "Ayurvedic skin care set",103    "Copper tongue scraper",104    "Cumin and ajwain detox water",105    "Handwritten Ayurvedic prescription",106    "Ayurvedic remedy instructions",107    "Doctor's note with herbal treatments",108    "Sanskrit Ayurvedic script",109    "Page from Ayurveda book",110    "Text describing herbal decoctions",111    "Prescription with cumin and ginger"112]113 114 115non_ayurvedic_prompts = [116    "Modern medicine bottle",117    "City hospital room",118    "Technology device on desk",119    "Fast food on a tray",120    "Traffic on an urban road",121    "Surgeon in operation theater",122    "Laptop and phone on a table",123    "Chemical pills in a box",124    "Modern doctor's office",125    "Processed food items",126    "Doctor using a stethoscope",127    "CT scan or X-ray room",128    "Cityscape with traffic",129    "Corporate office interior",130    "Plastic packaged medication",131    "Digital health monitoring device",132    "Blood pressure monitor",133    "Busy pharmacy counter",134    "Industrial lab equipment",135    "Person typing on a keyboard",136    "Microscope in a lab",137    "Person eating pizza",138    "High-tech surgical tools",139    "Injection or syringe closeup",140    "Coffee mug with laptop setup",141    "Capsules in a transparent strip",142    "Surgical gloves and mask",143    "Heart rate monitor screen",144    "Smartwatch fitness tracking",145    "Hospital bed with equipment",146    "Fast food burger meal",147    "IV drip setup",148    "Modern health supplement packaging",149    "White coat doctor with clipboard",150    "MRI machine in clinic",151    "Tech gadgets with health app",152    "Pharmaceutical advertisement",153    "Soda can with fries",154    "Crowded emergency room",155    "Modern pill organizer",156    "Fitness tracker close-up",157    "City skyline with pollution",158    "Desktop computer with reports",159    "Plastic bottles of vitamins",160    "Sterile lab setting",161    "Box of cough syrup",162    "Office desk with documents",163    "Ambulance in traffic",164    "Processed snack packs",165    "Online health consultation",166    "Medical chart on tablet",167    "Pills spilled on table",168    "Surgical face mask in hand",169    "Bottled energy drink",170    "X-ray images on display",171    "Blood sample in vial",172    "Digital thermometer reading",173    "Fitness coach with gadgets",174    "Disposable syringe packet",175    "Artificial sweeteners",176    "PowerPoint health presentation",177    "Surgery with robotic arms",178    "Skin care with chemical serum",179    "Health app interface",180    "Plastic pill blister pack",181    "Busy tech-driven clinic",182    "Athlete drinking sports drink",183    "Medical AI device demo",184    "Pharmacy shelf with labels",185    "High sugar soda bottle",186    "Biotech company office",187    "Keyboard with health analytics",188    "Doctor on a video call",189    "Smartphone showing medicine app",190    "Virtual health dashboard",191    "Injection for flu shot"192]193 194 195 196all_prompts = ayurvedic_prompts + non_ayurvedic_prompts197labels = ["Ayurveda"] * len(ayurvedic_prompts) + ["Non-Ayurveda"] * len(non_ayurvedic_prompts)198 199ayurvedic_keywords = ["herbs", "oil", "ayurveda", "healing", "natural", "plant", "therapy"]200 201def predict_with_clip(image_path, confidence_threshold=0.60):202    global clip_model, clip_preprocess203    if clip_model is None or clip_preprocess is None:204        clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)205 206    image = clip_preprocess(Image.open(image_path)).unsqueeze(0).to(device)207    text_inputs = clip.tokenize(all_prompts).to(device)208 209    with torch.no_grad():210        logits_per_image, _ = clip_model(image, text_inputs)211        probs = logits_per_image.softmax(dim=-1).cpu().numpy()[0]212 213    scores = {"Ayurveda": 0.0, "Non-Ayurveda": 0.0}214    for idx, prob in enumerate(probs):215        scores[labels[idx]] = max(scores[labels[idx]], prob)216 217    best_label = max(scores, key=scores.get)218    best_score = scores[best_label]219 220    if best_score >= confidence_threshold:221        return best_label, best_score, "CLIP"222    else:223        return None, best_score, "Low confidence"224 225def fallback_blip_classification(image_path):226    global blip_model, blip_processor227    if blip_model is None or blip_processor is None:228        from transformers import BlipProcessor, BlipForConditionalGeneration229        blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")230        blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)231 232    raw_image = Image.open(image_path).convert('RGB')233    inputs = blip_processor(raw_image, return_tensors="pt").to(device)234 235    with torch.no_grad():236        out = blip_model.generate(**inputs)237        caption = blip_processor.decode(out[0], skip_special_tokens=True)238 239    caption_lower = caption.lower()240    for kw in ayurvedic_keywords:241        if kw in caption_lower:242            return "Ayurveda", caption, "BLIP fallback"243 244    return "Non-Ayurveda", caption, "BLIP fallback"245 246def ocr_classification(image_path):247    text = pytesseract.image_to_string(Image.open(image_path))248    text_lower = text.lower()249 250    for kw in ayurvedic_keywords:251        if kw in text_lower:252            return "Ayurveda"253    return "Non-Ayurveda"254 255 256def classify_image(image_path):257    label = ocr_classification(image_path)258 259    if label == "Ayurveda":260        result = {261            "Type": label,262        }263    else:264        label, score, method = predict_with_clip(image_path)265        if method == "CLIP" and label is not None:266            result = {267                "Type": label,268            }269        else:270            label, caption, method = fallback_blip_classification(image_path)271            result = {272                "Type": label,273            }274 275    torch.cuda.empty_cache()276    gc.collect()277    return result