File size: 1,354 Bytes
e7840f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import os
from pathlib import Path
from PIL import Image
import pandas as pd
import Models
config = {
"model_root": "models",
"hf_model_repo": "SurfaceAI/models",
"models": {
"surface_type": "v1/surface_type_v1.pt",
"surface_quality": {
"asphalt": "v1/surface_quality_asphalt_v1.pt",
"concrete": "v1/surface_quality_concrete_v1.pt",
"paving_stones": "v1/surface_quality_paving_stones_v1.pt",
"sett": "v1/surface_quality_sett_v1.pt",
"unpaved": "v1/surface_quality_unpaved_v1.pt"
},
"road_type": "v1/road_type_v1.pt"
},
"gpu_kernel": 0,
"transform_surface": {
"resize": 384,
"crop": "lower_middle_half"
},
"transform_road_type": {
"resize": 384,
"crop": "lower_half"
},
}
root_path = Path(os.path.abspath(__file__)).parent
image_ids = [
"1351262795801113",
"153111940043147",
"1424818291203908",
]
image_data = []
for id in image_ids:
path = root_path / "example_images" / f"{id}.jpg"
try:
image_data.append(Image.open(path))
except Exception as e:
print(f'{e}: Not found or corrupted image: {path}')
continue
md = Models.ModelInterface(config=config)
df = md.batch_classifications(image_data, image_ids)
df.to_csv("example_prediction.csv")
|