Update app.py

app.py

@@ -8,148 +8,10 @@ from torch.utils.data import Dataset, DataLoader, ConcatDataset
 from PIL import Image, ImageTk # to display the image from the encoded pixels
 import gradio as gr
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel # importing the TrOCR processor representing the visual feature extrcator and tokenizer of the TrOCR model, and the TrOCR model
+import gradio as gr
 
 processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
-# loading and initializing with the pre-trained base trocr model
-model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten", low_cpu_mem_usage=True)
-# ADAM optimizer with decaying weights and the learning rate is set to 0.00005
-optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
-
-# configuring the model and setting undefined parameters
-# the decoder input_ids require the start and pad tokens, they are created by shifting the input to the right once
-model.config.decoder_start_token_id = processor.tokenizer.cls_token_id # id (=0) of the class token- <s> used as the first token after the inputs are shifted to the right for the decoder
-model.config.pad_token_id = processor.tokenizer.pad_token_id # id (=1) of the pad token- <pad>
-model.config.vocab_size = model.config.decoder.vocab_size # language modelling vocabulary size is set to default value of the decoder of the model (=50625)
-# sequence generation parameters associated with beam search (https://huggingface.co/blog/how-to-generate)
-model.config.eos_token_id = processor.tokenizer.sep_token_id # id (=2) of the separator token- </s> that is used at the end of the string, originally undefined
-model.config.max_length = 32 # maximum length to be used by the text generation function, originally 20
-model.config.num_beams = 5 # number of beams in beam search, beam width - number of sequences to consider while picking the one with the highest probablity, originally 1 (so only greedy search)
-model.config.early_stopping = True # beam search is stopped when 5(num_beams) sentences are done at a time (batch), originally False
-model.config.no_repeat_ngram_size = 3 # ngrams of size 3 can occur only once, to avoid word repetitions
-
-data_1 = pd.read_csv('washingtondb-v1.0/ground_truth/lines_truths.csv', header=None)
-images_location_1 = 'washingtondb-v1.0/data/line_images_normalized/'
-data_1.rename(columns={0: "file_name", 1: "text"}, inplace=True)
-data_1['file_name'] = [x + ".png" for x in data_1['file_name']]
-
-data_1['text'] = data_1['text'].str.replace('|',' ')
-# all characters in the ground truth have been separated by hyphens (-) so they are removed
-data_1['text'] = data_1['text'].str.replace('-','')
-
-data_1.drop(data_1.index[data_1['text'].str.contains('s_sl')], inplace = True)
-data_1.reset_index(drop=True, inplace=True)
-
-replace_chars = {'s_pt':'.', 's_qo':':','s_mi':'-', 's_bl':'(', 's_br':')', 's_sq':';', 's_s':'s', 's_cm':',', 's_et':'et', 's_lb':'£', 's_GW':'G.W.', 's_0':'0', 's_1':'1', 's_2':'2', 's_3':'3', 's_4':'4', 's_5':'5', 's_6':'6', 's_7':'7', 's_8':'8', 's_9':'9' }
-# iterating over the dictionary to replace the keys with the values
-for char in replace_chars.keys():
-  data_1['text'] = data_1['text'].str.replace(char, replace_chars[char])
-
-data_2 = pd.read_excel('Dates/Part I.xlsx')
-images_location_2 = 'Dates/Part I/'
-data_2 = data_2[:2]
-data_2.drop(['Image_Right', 'City', 'Category'], axis = 1, inplace = True) 
-data_2.rename(columns={"Image_Left": "file_name", "Date": "text"}, inplace=True)
-data_2 = data_2.astype({'text': 'str'})
-data_2['file_name'] = [x + ".jpg" for x in data_2['file_name']]
-
-train_1, val_1 = train_test_split(data_1, test_size=0.3) # splitting the data_1 dataframe into train and test sets with 80-20 split
-# train_1, val_1 = train_test_split(train_1, test_size=0.25) # further splitting the training set using 75-25 split for the validation set
-# the indices of the three data sets are reset in the same dataframe to freshly start from 0 each
-# instead of retaining indices from the original 'data' dataframe and the old indices are avoided being put as another column
-train_1.reset_index(drop=True, inplace=True)
-val_1.reset_index(drop=True, inplace=True)
-# test_1.reset_index(drop=True, inplace=True)
-train_2, val_2 = train_test_split(data_2, test_size=0.5) # splitting the data_2 dataframe into train and test datasets
-# train_2, val_2 = train_test_split(train_2, test_size=0.25) # further splitting for the validatio set
-# the indices of the three data sets are reset to 0 and old indices are avoided being put as another column
-train_2.reset_index(drop=True, inplace=True)
-val_2.reset_index(drop=True, inplace=True)
-# test_2.reset_index(drop=True, inplace=True)
-
-class ImageData(Dataset):
-    """
-    Class representing a custom PyTorch Dataset implementation with the images and their labels
-    """
-    def __init__(self, data, location):
-        """
-        Initialization Function 
-        """
-        self.data = data
-        self.location = location
-        self.processor = processor     
-
-    def __len__(self):
-        """
-        Function to get the number of samples in the Dataset
-        """
-        return len(self.data)
-
-    def __getitem__(self, idx):
-        """
-        Function/Getter for the contents of a WordData object at index- idx
-        Parameter:
-        idx - index at which the contents are to be retrieved
-        """
-        # to get the image file's name
-        img_file = self.location + self.data['file_name'][idx]
-        # resizing and normalizing the image to 3, 384, 384 (channels, image width and height) after removing the unnecessary 1 dimension bu using squeeze()
-        pixels = (self.processor(Image.open(img_file).convert("RGB"), return_tensors="pt").pixel_values).squeeze()
-        # encoding the text and getting the input ids or the encoded ground truth values 
-        # and all values until after the last value 128th one will be made the padding token
-        enc_values = self.processor.tokenizer(self.data['text'][idx], 
-                                          padding="max_length", 
-                                          max_length=32).input_ids
-
-        # encoded ground truth values are made into a tensor to use in computation 
-        # while the pad tokens each are set to -100 to be ignored during the computation of the loss
-        enc_values = torch.tensor([value if value != self.processor.tokenizer.pad_token_id else -100 for value in enc_values])
-
-        encodings = {"pixel_values": pixels, "labels": enc_values}
-        return encodings
-
-# Creating ImageDataset objects for training, validation and testing sets
-train_df_1 = ImageData(data=train_1, location=images_location_1)
-val_df_1 = ImageData(data=val_1, location=images_location_1)
-# test_df_1 = ImageData(data=test_1, location=images_location_1)
-
-train_df_2 = ImageData(data=train_2, location=images_location_2)
-val_df_2 = ImageData(data=val_2, location=images_location_2)
-# test_df_2 = ImageData(data=test_2, location=images_location_2)
-
-train_df = ConcatDataset([train_df_1,train_df_2])
-val_df = ConcatDataset([val_df_1,val_df_2])
-# test_df = ConcatDataset([test_df_1, test_df_2])
-
-# Creating DataLoaders for training, validation and testing data sets, each element is a batch of size 32
-# with data samples and all batches are shuffled at every epoch 
-train_dataloader = DataLoader(train_df, batch_size=32, shuffle=True)
-val_dataloader = DataLoader(val_df, batch_size=32, shuffle=True)
-# test_dataloader = DataLoader(test_df, batch_size=32, shuffle=True)
-
-epochs = 2
-losses = []
-
-# Training the model on the training data and validating using the validation set over epochs
-for epoch in range(epochs):
-   loss_sum = 0.0 # represents the sum of the loss loss computed after forward propagation of the inputs over each batch
-   model.train() # to put in training mode to activate Dropout and BatchNorm layers
-  
-  # iterating all the batches of data in the DataLoader with training data 
-  # and the pixel values and associated ground truth, keys and values of the WordData objects
-  # tqdm is used for displaying the progress bar covering each of the batches
-   for batch in val_dataloader:
-      # data is used on the same device as the model
-      # for key in batch.keys(): # iterating over the keys of the batch dictionary
-      #   batch[key] = batch[key].to(device)
-      optimizer.zero_grad() # resetting the gradients of optimized tensor to 0
-      outputs = model(**batch) # the input is passed to the model for forward propagation
-      loss = outputs.loss # represents the loss computed after forward propagation of the inputs
-      loss.backward() # backward propagation, calculating gradients
-      loss_sum += loss.item() # loss is converted into a number and is on the CPU now
-      optimizer.step() # optimizer is run, weights are updated
-      final_loss = loss_sum/len(train_dataloader)
-      losses.append(final_loss)
-
+model = VisionEncoderDecoderModel.from_pretrained("sk2003/hist-trocr")
 def process_image(image):
     # prepare image
     pixel_values = processor(image, return_tensors="pt").pixel_values
@@ -165,8 +27,8 @@ def process_image(image):
 title = "Hist-TrOCR"
 description = "Interactive demo of Hist-TrOCR, a fine-tuned version of Microsoft's TrOCR which is an end-to-end transformer model used for recognition of text from single-line or word images. It has been fine-tuned on historical text images. Upload an image or use one of the sample images below and click 'submit' to get the transcriptions. Results may take a few seconds to show up."
 
-# image1 = Image.open(images_location_1 + data_1['file_name'][10])
-# image = Image.open(images_location_1 + data_1['file_name'][11])
+image1 = Image.open(images_location_1 + data_1['file_name'][10])
+image = Image.open(images_location_1 + data_1['file_name'][11])
 # examples =[image1, image]
 
 iface = gr.Interface(fn=process_image, 
@@ -174,4 +36,4 @@ iface = gr.Interface(fn=process_image,
                      outputs=gr.outputs.Textbox(),
                      title=title,
                      description=description)
-iface.launch(debug=True)
+iface.launch(debug=True, share=True)