81 lines
2.3 KiB
Python
81 lines
2.3 KiB
Python
import os
|
|
import re
|
|
from typing import List
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw
|
|
|
|
|
|
class AnnotationRect:
|
|
"""Exercise 3.1"""
|
|
|
|
def __init__(self, x1, y1, x2, y2):
|
|
self.x1 = x1
|
|
self.x2 = x2
|
|
self.y1 = y1
|
|
self.y2 = y2
|
|
|
|
def area(self):
|
|
return (self.x2 - self.x1) * (self.y2 - self.y1)
|
|
|
|
def __array__(self) -> np.ndarray:
|
|
return np.array([self.x1, self.y1, self.x2, self.y2])
|
|
|
|
@staticmethod
|
|
def fromarray(arr: np.ndarray):
|
|
return AnnotationRect(arr[0], arr[1], arr[2], arr[3])
|
|
|
|
|
|
def read_groundtruth_file(path: str) -> List[AnnotationRect]:
|
|
"""Exercise 3.1b"""
|
|
annotationRects = []
|
|
with open(path, 'r') as file:
|
|
for line in file:
|
|
if line.strip():
|
|
values = line.strip().split()
|
|
annotationRects.append(AnnotationRect(float(values[0]), float(
|
|
values[1]), float(values[2]), float(values[3])))
|
|
return annotationRects
|
|
|
|
|
|
def get_image_with_max_annotations(dir_path: str) -> str:
|
|
img_pattern = re.compile(r'^(\d+)\.jpg$')
|
|
files = set(os.listdir(dir_path))
|
|
max_file = None
|
|
max_annotations = 0
|
|
|
|
for fname in files:
|
|
match = img_pattern.match(fname)
|
|
if match:
|
|
img_file = os.path.join(dir_path, fname)
|
|
annotations_number = len(read_groundtruth_file(os.path.join(
|
|
dir_path, f"{match.group(1)}.gt_data.txt")))
|
|
if (annotations_number > max_annotations):
|
|
max_file = img_file
|
|
max_annotations = annotations_number
|
|
return max_file
|
|
|
|
|
|
def visualize_image(image_path: str, output_path='output.jpg', rect_color=(255, 0, 0), width=2):
|
|
img_pattern = re.compile(r'(.*)(\.jpg)')
|
|
match = img_pattern.match(image_path)
|
|
annotations = read_groundtruth_file(f"{match.group(1)}.gt_data.txt")
|
|
|
|
img = Image.open(image_path).convert('RGB')
|
|
draw = ImageDraw.Draw(img)
|
|
|
|
for annotation in annotations:
|
|
draw.rectangle([annotation.x1, annotation.y1, annotation.x2, annotation.y2],
|
|
outline=rect_color, width=width)
|
|
|
|
img.save(output_path)
|
|
|
|
|
|
def main():
|
|
image_file = get_image_with_max_annotations(
|
|
"/home/ubuntu/mmp_wise2526_franksim/.data/mmp-public-3.2/train")
|
|
visualize_image(image_file)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|