Spaces:
Runtime error
Runtime error
| import torch | |
| import re | |
| import gradio as gr | |
| from pathlib import Path | |
| from transformers import AutoTokenizer, AutoFeatureExtractor, VisionEncoderDecoderModel | |
| # Pattern to ignore all the text after 2 or more full stops | |
| regex_pattern = "[.]{2,}" | |
| #sample = val_dataset[800] | |
| #model = model.cuda() | |
| #print(tokenizer.decode(model.generate(sample['pixel_values'].unsqueeze(0).cuda())[0]).replace('<|endoftext|>', '').split('\n')[0],'\n\n\n') | |
| def post_process(text): | |
| try: | |
| text = text.strip() | |
| text = re.split(regex_pattern, text)[0] | |
| except Exception as e: | |
| print(e) | |
| pass | |
| return text | |
| def predict(image, max_length=64, num_beams=4): | |
| image = image.convert('RGB') | |
| pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values | |
| pixel_values = pixel_values.to(device) | |
| with torch.no_grad(): | |
| text = tokenizer.decode(model.generate(pixel_values.cpu())[0]) | |
| text = text.replace('<|endoftext|>', '').split('\n') | |
| #[0],'\n\n\n' | |
| #text[0] | |
| #text = model.generate(pixel_values.cpu()) | |
| #text = tokenizer.decode(text.replace('<|endoftext|>', '').split('\n')[0],'\n\n\n') | |
| # output_ids = model.generate( | |
| # pixel_values, | |
| # max_length=max_length, | |
| # num_beams=num_beams, | |
| # return_dict_in_generate=True, | |
| #).sequences | |
| #preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True) | |
| #pred = post_process(preds[0]) | |
| return text[0] | |
| model_path = "team-indain-image-caption/hindi-image-captioning" | |
| device = "cpu" | |
| # Load model. | |
| model = VisionEncoderDecoderModel.from_pretrained(model_path) | |
| model.to(device) | |
| print("Loaded model") | |
| feature_extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k") | |
| print("Loaded feature_extractor") | |
| tokenizer = AutoTokenizer.from_pretrained(model_path) | |
| #if model.decoder.name_or_path == "gpt2": | |
| # tokenizer.pad_token = tokenizer.bos_token | |
| print("Loaded tokenizer") | |
| title = "Hindi Image Captioning" | |
| description = "" | |
| input = gr.inputs.Image(type="pil") | |
| #example_images = sorted([f.as_posix() for f in Path("examples").glob("*.jpg")]) | |
| #print(f"Loaded {len(example_images)} example images") | |
| interface = gr.Interface( | |
| fn=predict, | |
| inputs=input, | |
| outputs="textbox", | |
| title=title, | |
| description=description, | |
| #examples=example_images, | |
| live=True, | |
| ) | |
| interface.launch() | |