Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| from transformers import AutoFeatureExtractor, AutoTokenizer, VisionEncoderDecoderModel, GPT2Tokenizer, pipeline | |
| import os | |
| HF_DATASETS_OFFLINE=1 | |
| TRANSFORMERS_OFFLINE=1 | |
| device = 'cpu' | |
| auth_token = os.getenv("auth_token") | |
| #auth_token = os.environ.get("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") | |
| 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", use_auth_token=auth_token) | |
| tokenizer = GPT2Tokenizer.from_pretrained("gpt2", use_auth_token=auth_token) | |
| 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) | |
| 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 XRaySwinGen - Automatic Medical Report") | |
| if __name__ == "__main__": | |
| demo.launch() |