app.py

app.py

@@ -0,0 +1,96 @@
+import gradio as gr
+import torch
+from transformers import AutoFeatureExtractor, AutoTokenizer, VisionEncoderDecoderModel, GPT2Tokenizer, pipeline
+import os
+
+device = 'cpu'
+access_token = os.getenv("auth_token")
+
+max_length = 100
+num_beams = 4
+gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
+def predict_step(image_paths, model):
+  images = []
+  for image_path in image_paths:
+    i_image = Image.open(image_path)
+    if i_image.mode != "RGB":
+      i_image = i_image.convert(mode="RGB")
+    
+    #i_image.resize((640, 480))
+
+    images.append(i_image)
+
+  pixel_values = feature_extractor(images=images, return_tensors="pt").pixel_values
+  pixel_values = pixel_values.to(device)
+
+  output_ids = model.generate(pixel_values, **gen_kwargs)
+
+  preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+  preds = [pred.strip() for pred in preds]
+  return preds
+
+def predict_step_image(dataset_images, feature_extractor, model):
+  results = []
+  for i in dataset_images:
+    pixel_values = feature_extractor(images=i, return_tensors="pt").pixel_values
+    pixel_values = pixel_values.to(device)
+
+    output_ids = model.generate(pixel_values, **gen_kwargs)
+
+    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+    preds = [pred.strip() for pred in preds]
+    results.append(preds)
+  return results
+ 
+def predict_step_single_image(image, tokenizer, feature_extractor, model):
+    results=[]
+    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
+    pixel_values = pixel_values.to(device)
+
+    output_ids = model.generate(pixel_values, **gen_kwargs)
+
+    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+    preds = [pred.strip() for pred in preds]
+    results.append(preds)
+    return results
+
+def predict_step_pixel(dataset_pixel_values, model):
+  results=[]
+  for pv in dataset_pixel_values:
+    pixel_values = pv.reshape([1,3,224,224])
+    pixel_values = pixel_values.to(device)
+    output_ids = model.generate(pixel_values, **gen_kwargs)
+    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+    results.append([pred.strip() for pred in preds][0])
+  return results
+
+"""
+    image methods
+"""
+def load_image2txt_model(image_model_name):
+    model = VisionEncoderDecoderModel.from_pretrained(image_model_name)
+    feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/swin-large-patch4-window7-224")
+    
+    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+    tokenizer.pad_token = tokenizer.eos_token
+
+    model = model.to(device)
+    return tokenizer, feature_extractor, model
+
+def inference_image_pipe(image_input):
+    image_model_name = "./checkpoint-21000"
+
+    tokenizer, feature_extractor, image_model = load_image2txt_model(image_model_name)
+    #with autocast('cpu'):
+    text = predict_step_single_image(image_input, tokenizer, feature_extractor, image_model)[0]
+    return text
+
+with gr.Interface(fn=inference_image_pipe, 
+             inputs=gr.Image(shape=(256, 256)),
+             outputs="text",
+             examples=["3212210S4492629-1.png", "3216497S4499373-1.png"]) as demo:
+    gr.Markdown("POC V0 - XRay Automatic Medical Report")
+    
+    
+if __name__ == "__main__":
+    demo.launch(share=True)